import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler

# 设置中文字体和图表样式
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
sns.set_style("whitegrid")

# 加载波士顿房价数据集
print("正在加载数据...")
boston = fetch_openml(name='boston', version=1, as_frame=True)
df = boston.frame

# 显示数据基本信息
print(f"数据集形状: {df.shape}")
print(f"特征数量: {len(boston.feature_names)}")
print(f"目标变量: MEDV (房价中位数)")
print("\n前5行数据:")
print(df.head())

# 数据探索
print("\n=== 数据统计信息 ===")
print(df.describe())

print("\n=== 目标变量分布 ===")
print(df['MEDV'].describe())

# 特征与目标变量的相关性分析
correlation_matrix = df.corr()
print("\n=== 相关性矩阵 (前10个特征与房价的相关性) ===")
price_corr = correlation_matrix['MEDV'].sort_values(ascending=False)
print(price_corr)

# 可视化
fig, axes = plt.subplots(2, 3, figsize=(18, 12))
fig.suptitle('波士顿房价数据探索性分析', fontsize=16, fontweight='bold')

# 1. 房价分布
axes[0, 0].hist(df['MEDV'], bins=30, alpha=0.7, color='skyblue', edgecolor='black')
axes[0, 0].set_title('房价分布直方图')
axes[0, 0].set_xlabel('房价中位数 (MEDV)')
axes[0, 0].set_ylabel('频数')
axes[0, 0].axvline(df['MEDV'].mean(), color='red', linestyle='--', label=f'平均值: {df["MEDV"].mean():.2f}')
axes[0, 0].legend()

# 2. 相关性热力图
top_corr_features = price_corr.abs().nlargest(10).index
top_corr_matrix = df[top_corr_features].corr()
im = axes[0, 1].imshow(top_corr_matrix, cmap='coolwarm', aspect='auto')
axes[0, 1].set_xticks(range(len(top_corr_features)))
axes[0, 1].set_yticks(range(len(top_corr_features)))
axes[0, 1].set_xticklabels(top_corr_features, rotation=45, ha='right')
axes[0, 1].set_yticklabels(top_corr_features)
axes[0, 1].set_title('Top 10特征相关性热力图')
for i in range(len(top_corr_features)):
    for j in range(len(top_corr_features)):
        axes[0, 1].text(j, i, f'{top_corr_matrix.iloc[i, j]:.2f}',
                       ha='center', va='center', color='black', fontsize=8)

# 3. 房价与LSTAT的关系
axes[0, 2].scatter(df['LSTAT'], df['MEDV'], alpha=0.6, color='green')
axes[0, 2].set_title('房价 vs 低收入人群比例 (LSTAT)')
axes[0, 2].set_xlabel('低收入人群比例 (%)')
axes[0, 2].set_ylabel('房价中位数 (MEDV)')
z = np.polyfit(df['LSTAT'], df['MEDV'], 1)
p = np.poly1d(z)
axes[0, 2].plot(df['LSTAT'], p(df['LSTAT']), "r--", alpha=0.8)

# 4. 房价与RM的关系
axes[1, 0].scatter(df['RM'], df['MEDV'], alpha=0.6, color='orange')
axes[1, 0].set_title('房价 vs 平均房间数 (RM)')
axes[1, 0].set_xlabel('平均房间数')
axes[1, 0].set_ylabel('房价中位数 (MEDV)')
z = np.polyfit(df['RM'], df['MEDV'], 1)
p = np.poly1d(z)
axes[1, 0].plot(df['RM'], p(df['RM']), "r--", alpha=0.8)

# 5. 房价与PTRATIO的关系
axes[1, 1].scatter(df['PTRATIO'], df['MEDV'], alpha=0.6, color='purple')
axes[1, 1].set_title('房价 vs 师生比例 (PTRATIO)')
axes[1, 1].set_xlabel('师生比例')
axes[1, 1].set_ylabel('房价中位数 (MEDV)')
z = np.polyfit(df['PTRATIO'], df['MEDV'], 1)
p = np.poly1d(z)
axes[1, 1].plot(df['PTRATIO'], p(df['PTRATIO']), "r--", alpha=0.8)

# 6. 特征重要性预览（后续模型）
axes[1, 2].text(0.1, 0.8, '即将展示模型\n特征重要性', fontsize=14,
                transform=axes[1, 2].transAxes, ha='center',
                bbox=dict(boxstyle="round,pad=0.3", facecolor="lightblue", alpha=0.5))
axes[1, 2].set_title('模型特征重要性预览')
axes[1, 2].axis('off')

plt.tight_layout()
plt.show()

# 数据预处理
X = df.drop('MEDV', axis=1)
y = df['MEDV']

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 标准化特征
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# 训练线性回归模型
print("\n=== 线性回归模型 ===")
lr_model = LinearRegression()
lr_model.fit(X_train_scaled, y_train)

# 预测
y_pred_lr = lr_model.predict(X_test_scaled)

# 评估
mse_lr = mean_squared_error(y_test, y_pred_lr)
rmse_lr = np.sqrt(mse_lr)
r2_lr = r2_score(y_test, y_pred_lr)

print(f"线性回归模型性能:")
print(f"均方误差 (MSE): {mse_lr:.2f}")
print(f"均方根误差 (RMSE): {rmse_lr:.2f}")
print(f"R² 分数: {r2_lr:.4f}")

# 训练随机森林模型
print("\n=== 随机森林模型 ===")
rf_model = RandomForestRegressor(n_estimators=100, random_state=42)
rf_model.fit(X_train_scaled, y_train)

# 预测
y_pred_rf = rf_model.predict(X_test_scaled)

# 评估
mse_rf = mean_squared_error(y_test, y_pred_rf)
rmse_rf = np.sqrt(mse_rf)
r2_rf = r2_score(y_test, y_pred_rf)

print(f"随机森林模型性能:")
print(f"均方误差 (MSE): {mse_rf:.2f}")
print(f"均方根误差 (RMSE): {rmse_rf:.2f}")
print(f"R² 分数: {r2_rf:.4f}")

# 结果对比
results = pd.DataFrame({
    '模型': ['线性回归', '随机森林'],
    'MSE': [mse_lr, mse_rf],
    'RMSE': [rmse_lr, rmse_rf],
    'R²': [r2_lr, r2_rf]
})

print("\n=== 模型性能对比 ===")
print(results)

# 可视化预测结果
fig, axes = plt.subplots(2, 2, figsize=(16, 12))

# 1. 实际价格 vs 预测价格 (线性回归)
axes[0, 0].scatter(y_test, y_pred_lr, alpha=0.6, color='blue')
axes[0, 0].plot([y.min(), y.max()], [y.min(), y.max()], 'r--', lw=2)
axes[0, 0].set_xlabel('实际房价')
axes[0, 0].set_ylabel('预测房价')
axes[0, 0].set_title(f'线性回归: 实际 vs 预测\n(R² = {r2_lr:.3f})')
axes[0, 0].grid(True, alpha=0.3)

# 2. 实际价格 vs 预测价格 (随机森林)
axes[0, 1].scatter(y_test, y_pred_rf, alpha=0.6, color='green')
axes[0, 1].plot([y.min(), y.max()], [y.min(), y.max()], 'r--', lw=2)
axes[0, 1].set_xlabel('实际房价')
axes[0, 1].set_ylabel('预测房价')
axes[0, 1].set_title(f'随机森林: 实际 vs 预测\n(R² = {r2_rf:.3f})')
axes[0, 1].grid(True, alpha=0.3)

# 3. 残差分析 (线性回归)
residuals_lr = y_test - y_pred_lr
axes[1, 0].scatter(y_pred_lr, residuals_lr, alpha=0.6, color='orange')
axes[1, 0].axhline(y=0, color='r', linestyle='--')
axes[1, 0].set_xlabel('预测房价')
axes[1, 0].set_ylabel('残差')
axes[1, 0].set_title('线性回归: 残差分析')
axes[1, 0].grid(True, alpha=0.3)

# 4. 残差分析 (随机森林)
residuals_rf = y_test - y_pred_rf
axes[1, 1].scatter(y_pred_rf, residuals_rf, alpha=0.6, color='purple')
axes[1, 1].axhline(y=0, color='r', linestyle='--')
axes[1, 1].set_xlabel('预测房价')
axes[1, 1].set_ylabel('残差')
axes[1, 1].set_title('随机森林: 残差分析')
axes[1, 1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# 特征重要性 (随机森林)
feature_importance = pd.DataFrame({
    '特征': boston.feature_names,
    '重要性': rf_model.feature_importances_
}).sort_values('重要性', ascending=False)

print("\n=== 特征重要性排名 (随机森林) ===")
print(feature_importance)

# 可视化特征重要性
plt.figure(figsize=(12, 8))
bars = plt.barh(range(len(feature_importance)), feature_importance['重要性'],
                color='steelblue', alpha=0.7)
plt.yticks(range(len(feature_importance)), feature_importance['特征'])
plt.xlabel('特征重要性')
plt.title('波士顿房价预测 - 随机森林特征重要性')
plt.grid(axis='x', alpha=0.3)

# 在条形图上添加数值标签
for i, bar in enumerate(bars):
    width = bar.get_width()
    plt.text(width + 0.001, bar.get_y() + bar.get_height()/2,
             f'{width:.3f}', ha='left', va='center', fontsize=9)

plt.tight_layout()
plt.show()

# 预测示例
sample_idx = 0
sample_features = X_test.iloc[sample_idx:sample_idx+1]
sample_actual = y_test.iloc[sample_idx]
sample_pred_lr = lr_model.predict(scaler.transform(sample_features))[0]
sample_pred_rf = rf_model.predict(scaler.transform(sample_features))[0]

print(f"\n=== 预测示例 (第{sample_idx+1}个测试样本) ===")
print(f"实际房价: ${sample_actual*1000:.0f}")
print(f"线性回归预测: ${sample_pred_lr*1000:.0f}")
print(f"随机森林预测: ${sample_pred_rf*1000:.0f}")
print(f"特征值: {sample_features.iloc[0].to_dict()}")